import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

# 读取CSV文件
df = pd.read_csv('dataset.csv')

# 显示数据集基本信息
print("数据集形状:", df.shape)
print("\n前5行数据:")
print(df.head())

# 分离特征和目标变量
# 假设最后一列是y值，前面的列是x1到x9
X = df.iloc[:, :-1]  # 所有列除了最后一列
y = df.iloc[:, -1]   # 最后一列

print(f"\n特征矩阵形状: {X.shape}")
print(f"目标变量形状: {y.shape}")

# 按8:2的比例划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(
    X, y, 
    test_size=0.2,      # 20%作为验证集
    random_state=42,    # 设置随机种子以确保结果可重现
    shuffle=True        # 打乱数据
)

print(f"\n训练集特征形状: {X_train.shape}")
print(f"训练集目标变量形状: {y_train.shape}")
print(f"验证集特征形状: {X_val.shape}")
print(f"验证集目标变量形状: {y_val.shape}")

# 将训练集和验证集保存为CSV文件
# 训练集
train_data = pd.concat([X_train, y_train], axis=1)
train_data.to_csv('train_dataset.csv', index=False)

# 验证集
val_data = pd.concat([X_val, y_val], axis=1)
val_data.to_csv('val_dataset.csv', index=False)

print("\n数据集已成功划分并保存:")
print("- 训练集保存为: train_dataset.csv")
print("- 验证集保存为: val_dataset.csv")

# 显示划分后的统计信息
print(f"\n训练集样本数: {len(train_data)} ({len(train_data)/len(df)*100:.1f}%)")
print(f"验证集样本数: {len(val_data)} ({len(val_data)/len(df)*100:.1f}%)")

# 显示训练集和验证集的前几行
print("\n训练集前3行:")
print(train_data.head(3))

print("\n验证集前3行:")
print(val_data.head(3))